""" Module containing functions to read and plot aaomega spectra. """
from io import readtxt,writetxt
import numpy as np
import astro.spec
import astro.coord
import matplotlib.pyplot as pl
from matplotlib.mlab import rec_append_fields,rec_drop_fields
import pyfits

def read_fld(filename,skip=4):
    """ Read the info from an fld file."""
    fh = open(filename)
    f = fh.readlines()
    fh.close()
    f = [r for r in f if not r.lstrip().startswith('*')]
    radec = []
    for i,row in enumerate(f[skip:]):
        #print i
        row = row.split()
        ra,dec = ' '.join(row[1:4]), ' '.join(row[4:7])
        #print ra,dec
        radec.append(astro.coord.s2dec(ra,dec))

    fld = readtxt(f[skip:], names='name,ch,priority,mag,num,note',
                  usecols=(0,7,8,9,10,11))
    ra,dec = zip(*radec)
    fld = rec_append_fields(fld,'ra',ra)
    fld = rec_append_fields(fld,'dec',dec) 
    return fld

def readspec(filename):
    """Read AAOmega spectra from fits files.

    Returns a list of the spectra and a record array of the fibre info.
    """
    f = pyfits.open(filename)
    flux = f[0].data
    variance = f[1].data
    h = f[0].header

    spectra = [astro.spec.Spectrum(CRVAL=h['CRVAL1'],CDELT=h['CDELT1'],
                                   CRPIX=h['CRPIX1'],fl=fl,er=np.sqrt(var))
               for fl,var in zip(flux,variance)]

    for hdu in f:
        if 'MORE.FIBRES' in hdu.name:
            fibres = hdu.data
            break

    # convert ra, dec to degrees from radians
    for i in range(len(fibres)):
        fibres.RA[i] = fibres.RA[i] * 180. / np.pi
        fibres.DEC[i] = fibres.DEC[i] * 180. / np.pi

    return spectra,fibres

def plotobject(spectra, info, ax=None, comments=None,
               tweak_ymax=1.8, wmin=3700, wmax=8900, scale=True):
    """ Plots a series of aaomega spectra for a single object along
    with their info.
    """
    if ax is None:
        ax = pl.gca()
    ax.cla()
    a = ax
    cmap = 'brgmk'*(len(spectra) / 5 + 1)
    norm = 1
    ymax = ymin = 0
    alpha = 0.7
    sp = spectra[0]
    if sp is not None:
        fl = sp.fl[~np.isnan(sp.fl)]
        er = sp.er[~np.isnan(sp.er)]
        medfl = np.median(fl)
        sn = medfl / np.median(er) 
        norm = medfl
        sp.plot(ax=a, linewidth=0.5, alpha=alpha, yperc=0.85,
                flcolor=cmap[0], ercolor=cmap[0], show=False)
        y0,y1 = a.get_ylim()
        if (ymax == 0 or sn > 0.5):
            ymax = max(ymax, tweak_ymax*y1)
            ymin = min(ymin, y0)
        t = info[0]
        a.text((sp.wa.mean()-wmin)/(wmax-wmin), 0.9, '%s' % t['PIVOT'],
               transform=a.transAxes)

    for i,sp in enumerate(spectra[1:]):
        if sp is None:  continue
        fl = sp.fl[~np.isnan(sp.fl)]
        er = sp.er[~np.isnan(sp.er)]
        medfl = np.median(fl)
        sn = medfl / np.median(er)
        if scale:
            scalef = norm / medfl
            if scalef > 0:  sp.multiply(scalef)
        sp.plot(ax=a, linewidth=0.5, alpha=alpha, yperc=0.85,
                flcolor=cmap[i+1], ercolor=cmap[i+1], show=False)
        if scale:
            if scalef > 0:  sp.multiply(1./scalef)
        y0,y1 = a.get_ylim()
        if ymax == 0 or sn > 0.5:
            ymax = max(ymax, tweak_ymax*y1)
            ymin = min(ymin, y0)
        t = info[i+1]
        a.text((sp.wa.mean()-wmin)/(wmax-wmin), 0.9, '%s' % t['PIVOT'],
               transform=a.transAxes)

    a.set_xlim(wmin,wmax)
    a.set_ylim(ymin,ymax)
    a.set_autoscale_on(False)
    temp = tuple(t[k] for k in 'NAME RA DEC MAGNITUDE COMMENT'.split())
    a.text(0.16, 0.4, '%s\nRA,Dec %7.4f %6.4f\nMag %s\n%s\n' % temp,
           transform=a.transAxes)
    if comments is not None:
        a.text(0.7,0.5,comments[i],transform=a.transAxes)
    return

def plotconfig(spectra, info, outdir=None, indices=None,
               comments=None, tweak_ymax=1.8, wmin=3700, wmax=8900):
    """Plots fibres in a single configuration.

    Spectra and info are lists of spectra tuples, you can use this to
    plot more than one wavelength setting for a given spectrum.

    can also give a subset of fibres to plot (indices=), or save plots
    to a directory (outdir=). """
    spectra = zip(*spectra)
    info = zip(*info)
    assert len(spectra) == len(info), "Different lengths for spectra and info!"
    assert hasattr(spectra[0],'__iter__'), "Each spec entry should be a list"
    if indices is None:
        indices = range(len(spectra))

    if fig is None:
        fig = pl.figure()#figsize=(14, 9.9))
    else:
        fig.clf()

    fig.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95,
                        wspace=0.2, hspace=0.2)

    for i,index in enumerate(indices):
        if i % 4 == 0:
            pl.clf()
        a = fig.add_subplot(4,1,i % 4 + 1)
        plotobject(spectra[index], info[index], ax=a, comments=None,
                   tweak_ymax=tweak_ymax, wmin=wmin, wmax=wmax)
